from os.path import join, abspath
import json
from backend.experiment.clue.afqmc.scheme0.implement \
    import AFQMCImplement as S0Impl
from typing import List


SAMPLE_LIMIT = 50


class AFQMCImplement(S0Impl):

    @classmethod
    def _read_train_records(cls) -> list:
        with open(join(cls.data_dir(), 'train.json'), 'r', encoding='utf8') \
                as file:
            records = [
                json.loads(each) for each in file.readlines()[:SAMPLE_LIMIT]]
        return records

    @classmethod
    def label_map(cls) -> List[str]:
        return ['不同', '相似']

    @classmethod
    def _make_prompt(cls, record, mask_target: bool) -> str:
        s1: str = record['sentence1']
        s2: str = record['sentence2']
        label: str = cls.__mask_token if mask_target else \
            cls.label_map()[int(record['label'])]
        return f'下面两个句子语义{label}：{s1}。{s2}'

    @classmethod
    def cache_dir(cls) -> str:
        return join(cls.work_dir(), 'cache')

    @classmethod
    def work_dir(cls) -> str:
        return abspath(join(__file__, '..'))
